//
// Created by Administrator on 2021/7/28.
//

#include <vector>
#include <iostream>
#include <unordered_map>

using namespace std;

// Definition for a binary tree node.
struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;

    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
};

class Solution {
private:
    unordered_map<int, TreeNode *> father; // hash 各个节点对应的父节点
    vector<int> ans; // 储存答案
    // 搜索父节点
    void findFather(TreeNode *root) {
        if (root->left != nullptr) {
            father[root->left->val] = root;
            findFather(root->left);
        }
        if (root->right != nullptr) {
            father[root->right->val] = root;
            findFather(root->right);
        }
    }

    // 搜索答案
    void findAns(TreeNode *root, TreeNode *from, int range, int k) {
        if (root == nullptr) return;
        if (range == k) {
            this->ans.push_back(root->val);
            return;
        }
        // 通过判断是否和来源节点相同，减少判断次数
        if (root->left != from) {
            findAns(root->left, root, range + 1, k);
        }
        if (root->right != from) {
            findAns(root->right, root, range + 1, k);
        }
        if (father[root->val] != from) {
            findAns(father[root->val], root, range + 1, k);
        }
    }

public:
    vector<int> distanceK(TreeNode *root, TreeNode *target, int k) {
        // 先从根节点进行一次dfs，找到各个结点的父节点
        findFather(root);
        // 再从目标节点出发，搜索距离为k的
        findAns(target, nullptr, 0, k);
        return this->ans;
    }
};

int main() {
    auto n1 = TreeNode(3), n2 = TreeNode(5), n3 = TreeNode(1), n4 = TreeNode(6), n5 = TreeNode(2),
            n6 = TreeNode(0), n7 = TreeNode(8), n8 = TreeNode(7), n9 = TreeNode(4);
    n1.left = &n2;
    n1.right = &n3;
    n2.left = &n4;
    n2.right = &n5;
    n3.left = &n6;
    n3.right = &n7;
    n5.left = &n8;
    n5.right = &n9;
    Solution solution;
    auto ans = solution.distanceK(&n1, &n2, 2);
    for (auto &x:ans) cout << x << endl;
    return 0;
}